from __future__ import annotations
import torch
import torch.nn as nn

from abc import abstractmethod
from analytics.crop_recognition.evaluators.evaluator import Evaluator


class BaseClassificationModel(nn.Module):
    @abstractmethod
    def __init__(self):
        super().__init__()

    @abstractmethod
    def forward(self, src: torch.FloatTensor, src_padding_mask: torch.BoolTensor):
        raise NotImplementedError()

    @staticmethod
    @abstractmethod
    def update_evaluator(
        model: BaseClassificationModel,
        src: torch.FloatTensor,
        src_padding_mask: torch.BoolTensor,
        label: torch.LongTensor,
        label_weight: torch.FloatTensor,
        evaluator: Evaluator,
    ):
        raise NotImplementedError()

    @staticmethod
    @abstractmethod
    def calculate_loss(
        model: BaseClassificationModel,
        src: torch.FloatTensor,
        src_padding_mask: torch.BoolTensor,
        label: torch.LongTensor,
        label_weight: torch.FloatTensor,
        scale: float,
    ):
        raise NotImplementedError()
